package com.tuniu.agents.impl.trip;

import com.tuniu.agents.agent.Agent;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.agent.ToolAgent;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.tools.trip.tool.ItineraryPlanListTool;
import com.tuniu.agents.util.ChatResponseUtils;
import com.tuniu.agents.util.I18nUtil;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.*;

@Slf4j
public class Designer extends ToolAgent {

    public static final String AGENT_ID = "Designer";

    public static final String PROMPT_ID = "designer";

    protected final AgentManager agentManager;


    public Designer(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, String... tools) {
        super(agentManagerGroup, agentId, promptTemplateId, tools);
        this.agentManager = agentManagerGroup.getAgentManager();
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String,Object> context) {
        ChatResponse chatResponse = this.callResponse(conversationId, messages);
        String toolName = ChatResponseUtils.getFirstToolName(chatResponse);
        if (!StringUtils.hasText(toolName)) {
            return Flux.just(ChatResponseUtils.getFirstTextMessage(chatResponse));
        }else {
            String toolResponse = ChatResponseUtils.getFirstTextMessage(chatResponse);
            if (!StringUtils.hasText(toolResponse) || toolResponse.equals("null")) {
                log.error("Designer agent toolResponse fail:{}", toolResponse);
                return Flux.just(BaseAgentNewResp.sendAnswerMsg(null, I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY)));
            }

            if(toolName.equals(ItineraryPlanListTool.TOOLNAME)){

                log.info("Designer toolResponse:{}", toolResponse);
                ItineraryPlanListTool.Response res = JsonUtil.fromJson(toolResponse, ItineraryPlanListTool.Response.class);

                if(null != res && CollectionUtils.isNotEmpty(res.data())){

                    Set<String> scenicSpotsSet = new HashSet<>();
                    Set<String> hotelSet = new HashSet<>();
                    res.data().forEach(itineraryPlan -> {
                        addSplitItemsToSet(scenicSpotsSet, itineraryPlan.itineraryScenicSpots());
                        addSplitItemsToSet(hotelSet, itineraryPlan.itineraryHotels());
                    });

                    StringBuilder sb = new StringBuilder();
                    sb.append("\n");
                    if(CollectionUtils.isNotEmpty(scenicSpotsSet)){
                        sb.append("### scenic spots\n");
                        sb.append(String.join("、", scenicSpotsSet)).append("\n");
                    }
                    if(CollectionUtils.isNotEmpty(hotelSet)){
                        sb.append("### hotel\n");
                        sb.append(String.join("、", hotelSet)).append("\n");
                    }

                    Agent agent = agentManager.getAgent(Planner.AGENT_ID);
                    Map<String, Object> hashMap = chatEnvironmentMemory.get(conversationId);
                    hashMap.put("itineraryResources",sb.toString());
                    chatEnvironmentMemory.add(conversationId, hashMap);
                    return agent.stream(conversationId, messages);
                }

            }

        }

        return Flux.just(BaseAgentNewResp.sendAnswerMsg(null, I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY)));
    }

    private void addSplitItemsToSet(Set<String> set, String items) {
        if (StringUtils.hasText(items)) {
            set.addAll(Arrays.asList(items.split("、")));
        }
    }
}
